# Import libraries
import requests
from pathlib import Path
from tqdm import tqdm
import argparse


# Define the function to download the model
parser = argparse.ArgumentParser(description='Process some integers.')
parser.add_argument(
    '--model', 
    choices=['yolov10n', 'yolov10s', 'yolov10m', 'yolov10b', 'yolov10l', 'yolov10x'],
    default='yolov10n', 
    help='Model to download'
)


def download_model(model, save_path, max_retries=5):
    '''
        Function to download the model from the github release page
    '''

    url = "https://github.com/THU-MIG/yolov10/releases/download/v1.1/" + model + ".pt"
    # url = "https://hf-mirror.com/onnx-community/YOLOv10/blob/main/" + model +".onnx"

    save_path = Path(save_path)
    save_path.parent.mkdir(parents=True, exist_ok=True)
    
    for attempt in range(max_retries):
        try:
            response = requests.get(url, stream=True, timeout=10)
            response.raise_for_status()
            
            total_size = int(response.headers.get("content-length", 0))
            progress_bar = tqdm(
                desc=str(save_path),
                total=total_size,
                unit="B",
                unit_scale=True,
                unit_divisor=1024,
            )
            
            with open(save_path, "wb") as file:
                for chunk in response.iter_content(chunk_size=8192):
                    file.write(chunk)
                    progress_bar.update(len(chunk))
            progress_bar.close()
            return True
        except requests.exceptions.RequestException as e:
            print(f"错误：{e}(重试 {attempt+1}/{max_retries})")
    return False

if __name__ == "__main__":
    args = parser.parse_args()
    # download_model("yolov10n", "yolov10n.onnx")
    download_model("yolov10n", "yolov10n.pt")
